#include "sm3.h"
#include <stdio.h>
#include <string.h>

/**
 * SM3上下文初始化函数.
 *
 * \param ctx 上下文
 * \return 返回1表示正确
 */
int sm3_init(sm3_context *ctx)
{
	memset(ctx, 0, sizeof(*ctx));

	ctx->state[0] = SM3_A;
	ctx->state[1] = SM3_B;
	ctx->state[2] = SM3_C;
	ctx->state[3] = SM3_D;
	ctx->state[4] = SM3_E;
	ctx->state[5] = SM3_F;
	ctx->state[6] = SM3_G;
	ctx->state[7] = SM3_H;

	return 1;
}

/**
 * SM3的压缩函数。将Vi和Bi压缩成V(i+1).
 *
 * \param ctx 上下文
 * \param data Bi。消息填充并按512个比特进行分块。Bi表示其中的第i块。
 * \return 返回1表示正确
 */
int sm3_compress_function(sm3_context *ctx, const uint8_t data[64])
{
	// 变量定义
	// 消息扩展的第一部分和第二部分
	uint32_t W[68], W1[64];
	// A, B, C, D, E, F, G, H 8个字寄存器
	uint32_t A, B, C, D, E, F, G, H;
	// 中间变量
	uint32_t SS1, SS2, TT1, TT2;

	// 1.消息扩展
	// 将 512 个比特数据扩展成 67584 个比特（132个字）
	// 1.1.第一步 将消息分组Bi分成16个字
	for (int i = 0; i < 16; i++)
	{
		GET_UINT_BE(W[i], data, i * 4);
	}

	// 1.2.第二步 扩展W部分。范围[16,67]
	for (int i = 16; i < 68; i++)
	{
		// Wi = P1( Wi-16 ^ Wi-9􀱇^(Wi-3<<<15))􀱇^ (Wi-13<<<7)􀱇^ Wi-6
		// W[i] = P1(W[i - 16] ^ W[i - 9] ^ ROTATE(W[i - 3], 15)) ^ ROTATE(W[i - 13], 7) ^ W[i - 6];
		W[i] = EXPAND(W[i - 16], W[i - 9], W[i - 3], W[i - 13], W[i - 6]);

		// printf("%08X\t", W[i]);
	}

	// print_uint32_arr_by_hex("W0W1…W67",W, 68);

	// 1.3.第三步 扩展W'部分。范围[0,63]
	for (int i = 0; i < 64; i++)
	{
		W1[i] = W[i] ^ W[i + 4];
	}

	// print_uint32_arr_by_hex("W'0W'1…W'63", W1, 64);

	// 2.压缩
	// 2.1.取出Vi
	A = ctx->state[0];
	B = ctx->state[1];
	C = ctx->state[2];
	D = ctx->state[3];
	E = ctx->state[4];
	F = ctx->state[5];
	G = ctx->state[6];
	H = ctx->state[7];

	// 2.2.开始压缩
	// 2.2.1.使用循环压缩 压缩次数64次
	// 需要使用4个布尔函数压缩 前面16个次压缩使用布尔函数FF0和GG0，后面48次压缩使用FF1和GG1
	// printf("\t\t%8C\t%8C\t%8C\t%8C\t%8C\t%8C\t%8C\t%8C\n", 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H');
	// printf("\t\t%08X\t%08X\t%08X\t%08X\t%08X\t%08X\t%08X\t%08X\n",A, B, C, D, E, F, G, H);

	for (int j = 0; j < 64; j++)
	{

		// 需要前面 16 次压缩和后面 48 次压缩要分开处理
		// 2.2.2 SS1=((A<<<12)+E+(Ti<<<(jmod32)))<<<7
		// 2.2.3. SS2 = SS1 ^ (A<<<12)
		// 2.2.4. TT1 = FFi(A,B,C)+D+SS2+W'i
		// 2.2.4. TT2 = GGi(E,F,G)+H +SS1+Wi

		if (j < 16)
		{
			// 2.2.2 SS1=((A<<<12)+E+(Ti<<<(jmod32)))<<<7
			SS1 = ROTATE(ROTATE(A, 12) + E + ROTATE(SM3_CONST_T0, j % 32), 7);

			// 2.2.3. SS2 = SS1 ^ (A<<<12)
			SS2 = SS1 ^ ROTATE(A, 12);

			// 2.2.4. TT1 = FFi(A,B,C)+D+SS2+W'i
			TT1 = FF0(A, B, C) + D + SS2 + W1[j];

			// 2.2.4. TT2 = GGi(E,F,G)+H +SS1+Wi
			TT2 = GG0(E, F, G) + H + SS1 + W[j];
		}
		else
		{
			// 2.2.2 SS1=((A<<<12)+E+(Ti<<<(jmod32)))<<<7
			SS1 = ROTATE(ROTATE(A, 12) + E + ROTATE(SM3_CONST_T1, j % 32), 7);

			// 2.2.3. SS2 = SS1 ^ (A<<<12)
			SS2 = SS1 ^ ROTATE(A, 12);

			// 2.2.4. TT1 = FFi(A,B,C)+D+SS2+W'i
			TT1 = FF1(A, B, C) + D + SS2 + W1[j];

			// 2.2.4. TT2 = GGi(E,F,G)+H +SS1+Wi
			TT2 = GG1(E, F, G) + H + SS1 + W[j];
		}

		// 2.2.5. D = C
		D = C;

		// 2.2.6. C = B << < 9
		C = ROTATE(B, 9);

		// 2.2.7. B = A
		B = A;

		// 2.2.8. A = TT1
		A = TT1;

		// 2.2.9. H = G
		H = G;

		// 2.2.10 G = F << < 19
		G = ROTATE(F, 19);

		// 2.2.11 F = E
		F = E;

		// 2.2.12 E = P0(TT2)
		E = P0(TT2);

		// printf("%8d\t%08X\t%08X\t%08X\t%08X\t%08X\t%08X\t%08X\t%08X\n",j, A, B, C, D, E, F, G, H);
	}

	// 3.更新上下文
	ctx->state[0] ^= A;
	ctx->state[1] ^= B;
	ctx->state[2] ^= C;
	ctx->state[3] ^= D;
	ctx->state[4] ^= E;
	ctx->state[5] ^= F;
	ctx->state[6] ^= G;
	ctx->state[7] ^= H;

	return 1;
}

/**
 * 迭代函数.
 *
 * \param ctx 上下文
 * \param input 输入
 * \param input_len 输入长度
 * \return 返回1表示正确
 */
int sm3_update(sm3_context *ctx, uint8_t *input, const size_t input_len)
{
	size_t block_size = 64;
	size_t block_num = input_len / block_size;
	ctx->len = input_len;

	// 分两部分处理
	// 1.完整的块 一个块 512 个比特（64个字节）
	// 依次压缩
	for (size_t i = 0; i < block_num; i++)
	{
		memcpy(ctx->buffer, input, block_size);

		// 压缩
		sm3_compress_function(ctx, ctx->buffer);

		input += block_size;
	}

	// 2.不完整的块
	if (input_len > block_num * block_size)
	{
		// 将剩下的内容放到缓冲区里面
		// 剩下的内容的长度  input_len - block_num * block_size
		memcpy(ctx->buffer, input, input_len - block_num * block_size);
	}

	return 1;
}

/**
 * 处理最后一个块，并获取最终的摘要值。最后一个块 = 未处理完的消息 + 填充值 + 消息长度.
 *
 * \param output 输出
 * \param ctx 上下文
 * \return 返回1表示正确
 */
int sm3_finish(uint8_t output[32], sm3_context *ctx)
{
	// 处理最有一个块
	// 1.最后一个块里面的数据。因为最后一个块的数据已经在缓冲区里面，所以这里不用管
	size_t block_size = 64;
	size_t block_num = ctx->len / block_size;
	// 剩余数据的长度
	size_t remain_len = ctx->len - block_num * block_size;

	// 2.填充数据
	// 需要填充的长度 0-56
	size_t data_len = 56;
	size_t fill = data_len - remain_len;
	memcpy(ctx->buffer + remain_len, sm3_padding, fill);

	// 4.设置长度
	int size_len = 8;
	uint8_t input_len[8];
	// 将消息的长度的 整型 转 字符数组
	// 长度是单位是字节，这里需要转比特
	PUT_ULONG_BE(ctx->len * 8, input_len, 0);

	// 处理最后一个单元，将长度放进去
	memcpy(ctx->buffer + data_len, input_len, size_len);

	// 5.哈希
	sm3_compress_function(ctx, ctx->buffer);

	// 6.输出结果
	for (int i = 0; i < 8; i++)
	{
		PUT_UINT_BE(ctx->state[i], output, i * 4);
	};

	return 1;
}

/**
 * 使用16进制的格式打印字符数组.
 *
 * \param desc 说明
 * \param chars 字符数组
 * \param len 字符数组的长度
 */
void print_chars_by_hex(const int8_t desc[], const uint8_t chars[], const size_t len)
{
	printf("打印%s：\n", desc);

	for (size_t i = 0; i < len; i++)
	{
		printf("%02X", chars[i]);

		if ((i + 1) % 4 == 0)
		{
			printf("\t");
		}

		if ((i + 1) % 32 == 0)
		{
			printf("\n");
		}
	}

	printf("\n");
}

/**
 * 使用16进制的格式打印 uint32 数组.
 *
 * \param desc 说明
 * \param arr uint32_t 数组
 * \param len uint32_t 数组的长度
 */
void print_uint32_arr_by_hex(const int8_t desc[], const uint32_t arr[], const size_t len)
{
	printf("打印%s：\n", desc);

	uint8_t print_arr[4];

	for (size_t i = 0; i < len; i++)
	{
		PUT_UINT_BE(arr[i], print_arr, 0);

		for (size_t i = 0; i < 4; i++)
		{
			printf("%02X", print_arr[i]);
		}

		printf("\t");

		if ((i + 1) % 8 == 0)
		{
			printf("\n");
		}
	}

	printf("\n");
}
